#!/usr/bin/python
'''
Parses paired ham files.
Ham files must be sorted by name.
'''
# System imports.
import os
import sys
import string
import time
import random
import logging
import ham
import numpy as np

#logging.basicConfig(level=logging.DEBUG, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )
logging.basicConfig(level=logging.INFO, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )

## Parameters.
ham_file1 = sys.argv[1]
ham_file2 = sys.argv[2]
ham_out_file = sys.argv[3]
insert_size = int(sys.argv[4])
std_dev = int(sys.argv[5])
read_size = int(sys.argv[6])
key1 = sys.argv[7]
key2 = sys.argv[8]
orientation = sys.argv[9]
ctg_size_file = sys.argv[10]

chunk_size = ham.default_chunk_size
#chunk_size = 10

# read contig sizes.
logging.info("reading contig sizes.")
fin = open(ctg_size_file, "rb")
lines = fin.readlines()
fin.close()

ctg_size = {}
for line in lines:
	tmp = line.strip().split()
	ctg_size[tmp[1]] = int(tmp[0])


# setup intra-contig bounds
intra_min = insert_size - (5 * std_dev)
intra_max = insert_size + (5 * std_dev)

# Note size of end keys.
ksz1 = len(key1)
ksz2 = len(key2)

# Open pointers to ham objects.
ham_obj1 = ham.HamFile(ham_file1, read=1)
ham_obj2 = ham.HamFile(ham_file2, read=1)
sz1 = ham_obj1.data.len()
sz2 = ham_obj2.data.len()

# Allocate annotation array.
logging.info("creating annotation arrays of size %i, %i." % (sz1, sz2))
ant_arr1 = ham.create_annotate_array(sz1)
ant_arr2 = ham.create_annotate_array(sz2)

# load first chunk of data.
logging.info("loading first chunk of ham arrays.")
ham_idx_1 = 0
ham_idx_2 = 0
chunk_idx_1 = 0
chunk_idx_2 = 0
ham_arr1 = ham_obj1.data[ham_idx_1:ham_idx_1 + chunk_size, 'QNAME', 'FLAG', 'RNAME', 'POS']
ham_arr2 = ham_obj2.data[ham_idx_1:ham_idx_2 + chunk_size, 'QNAME', 'FLAG', 'RNAME', 'POS']
ham_idx_1 = ham_idx_1 + chunk_size
ham_idx_2 = ham_idx_2 + chunk_size
local_i = 0
local_j = 0

# find first valid entries.
logging.info("boot strapping interation.")
for x in range(sz1):
	if ham_arr1[x]['FLAG'][0] == 0 or ham_arr1[x]['FLAG'][0] == 16:
		break
for y in range(sz2):
	if ham_arr2[y]['FLAG'][0] == 0 or ham_arr2[y]['FLAG'][0] == 16:
		break

# bootstrap iteration.
buf1 = set()
buf2 = set()
if x < y:	
	pname1 = ham_arr1[x]['QNAME'][0][:-ksz1]
	pname2 = False
	bufname = pname1
	buf1.add((chunk_idx_1, local_i, x))
	i = x
	j = y - 1
else:
	pname2 = ham_arr2[y]['QNAME'][0][:-ksz2]
	pname1 = False
	bufname = pname2
	buf2.add((chunk_idx_2, local_j, y))
	i = x - 1
	j = y
k = 0

#logging.debug("%s\t%s\t%s" % (pname1, pname2, bufname))

# make output arrays.
logging.info("making output arrays.")
ham_arr1_perfect = ham.create_array(5000000)
ham_arr1_ok = ham.create_array(5000000)

ham_arr2_perfect = ham.create_array(5000000)
ham_arr2_ok = ham.create_array(5000000)

'''
Each iteration of this while loop increases one of i or j.
The one that is increased is simply the one that is 
lexicographically smaller than the other. This comparison
uses the numpy implementation.
'''
logging.info("traversing sorted ham arrays.")
oki = 0
peri = 0
while 1 == 1:

	# progress.
	'''
	if i % chunk_size == 0:
		logging.info("progress arr1 %i of %i" % (i, sz1))
	if j % chunk_size == 0:
		logging.info("progress %i of %i" % (j, sz2))
	'''
	
	# see who to change.
	if pname1 < pname2:
		# update index.
		i += 1
		local_i += 1
		
		# check for break.
		if i == sz1: break	
		
		# check if we need to load more for this.
		if local_i >= ham_arr1.size:
			# load a new chunk.
			prev_ham_arr1 = ham_arr1
			
			if ham_idx_1 + chunk_size > sz1:
				topidx = sz1
			else:
				topidx = ham_idx_1 + chunk_size
				
			if ham_idx_1 >= topidx:
				break
			
			logging.info("loading new chunk, %i of %i of ham1 array." % (topidx, sz1))
			ham_arr1 = ham_obj1.data[ham_idx_1:topidx, 'QNAME', 'FLAG', 'RNAME', 'POS']
			ham_idx_1 = topidx
			local_i = 0
			chunk_idx_1 += 1
			
			#if chunk_idx_1 >= 1: break
		
		# increment left.
		pname1 = ham_arr1[local_i]['QNAME'][0][:-ksz1]	# get name minus tag.
		choice = True					# save which name was updated.

	else:
		# update index.
		j += 1
		local_j += 1

		# check for break.
		if j == sz2: break

		# check if we need to load more for this.
		if local_j >= ham_arr2.size:
			# load a new chunk.
			logging.info("loading new chunk, %i of %i of ham2 array." % (ham_idx_2, sz2))
			prev_ham_arr2 = ham_arr2
			
			if ham_idx_2 + chunk_size > sz2:
				topidx = sz2
			else:
				topidx = ham_idx_2 + chunk_size

			if ham_idx_2 >= topidx:
				break
			
			ham_arr2 = ham_obj2.data[ham_idx_2:topidx, 'QNAME', 'FLAG', 'RNAME', 'POS']
			ham_idx_2 = topidx
			local_j = 0
			chunk_idx_2 += 1
		
		# increment right.
		pname2 = ham_arr2[local_j]['QNAME'][0][:-ksz1]
		choice = False
		
		
	# make partner index -1.
	ant_arr1[i]['MIDX'] = -1
	ant_arr2[j]['MIDX'] = -1
	
	
	#print sz1, sz2, i, j, pname1, pname2
	#logging.debug("%s\t%s\t%s\t%i\t%i" % (bufname, pname1, pname2, i, j))
		
	# skip non-mapped reads.
	if ham_arr1[local_i]['FLAG'][0] != 0 and ham_arr1[local_i]['FLAG'][0] != 16: continue
	if ham_arr2[local_j]['FLAG'][0] != 0 and ham_arr2[local_j]['FLAG'][0] != 16: continue
	
	# add line to buffer if matches.
	test1 = pname1 == bufname
	test2 = pname2 == bufname
	if test1 == True:
		buf1.add((chunk_idx_1, local_i, i))
	if test2 == True:
		buf2.add((chunk_idx_2, local_j, j))
		
	# if one name matchs buffer continue.
	if test1 == True or test2 == True: continue
	
	# process buffer.
	bsz1 = len(buf1)
	bsz2 = len(buf2)
	
	# 1: unique mapped	/	unmapped
	if bsz1 == 1 and bsz2 == 0:
		e = buf1.pop()
		ant_arr1[e[2]]['TYPE'] = 1
		logging.debug(1)
		
	# 1: unmapped	/	unique mapped
	elif bsz1 == 0 and bsz2 == 1:
		e = buf2.pop()
		ant_arr2[e[2]]['TYPE'] = 1
		logging.debug(1)
		
	# 2: unique mapped	/	unique mapped
	elif bsz1 == 1 and bsz2 == 1:
		e1 = buf1.pop()
		e2 = buf2.pop()
	
		
		if e1[0] == chunk_idx_1:
			rid1 = ham_arr1[e1[1]]['RNAME'][0]
			f1 = ham_arr1[e1[1]]['FLAG'][0]
			pos1 = ham_arr1[e1[1]]['POS'][0]
					
		else:
			rid1 = prev_ham_arr1[e1[1]]['RNAME'][0]
			f1 = prev_ham_arr1[e1[1]]['FLAG'][0]
			pos1 = prev_ham_arr1[e1[1]]['POS'][0]
			
			
		if e2[0] == chunk_idx_2:
			rid2 = ham_arr2[e2[1]]['RNAME'][0]
			f2 = ham_arr2[e2[1]]['FLAG'][0]
			pos2 = ham_arr2[e2[1]]['POS'][0]	
		else:
			rid2 = prev_ham_arr2[e2[1]]['RNAME'][0]
			f2 = prev_ham_arr2[e2[1]]['FLAG'][0]
			pos2 = prev_ham_arr2[e2[1]]['POS'][0]

		
		# check orientation constraints.
		opassed = True
		if rid1 == rid2:
			if orientation == "illumina":
				if f1 == f2:
					opassed = False
			elif orientation == "solid":
				if f1 != f2:
					opassed = False	
					
		# check for intra-contig insert size.
		icpassed = True
		if rid1 == rid2:
			diff = abs(pos1 - pos2)
			if diff < intra_min or diff > intra_max:
				icpassed = False
				
		
		# check for edge of contig insert sizes.
		ebpassed = True
		if rid1 != rid2:
			# find closest one edge.
			ltmp = ctg_size[rid1] - (pos1 + read_size)
			if pos1 < ltmp:
				edge1 = pos1
			else:
				edge1 = ltmp
				
			# find closest two edge.				
			ltmp = ctg_size[rid2] - (pos2 + read_size)
			if pos2 < ltmp:
				edge2 = pos2
			else:
				edge2 = ltmp
				
			# see if this bound is violated.
			if edge1 + edge2 > intra_max:
				ebpassed = False
				
					
		# save status.
		'''
		if opassed == False and icpassed == False:
			# mark as bad orientation and bad insert size.
			ant_arr1[e1[2]]['TYPE'] = 3
			ant_arr2[e2[2]]['TYPE'] = 3
			ant_arr1[e1[2]]['MIDX'] = e2[2]
			ant_arr2[e2[2]]['MIDX'] = e1[2]			
			
		elif opassed == False and icpassed == True:
			# mark as bad orientation and good insert size.
			ant_arr1[e1[2]]['TYPE'] = 4
			ant_arr2[e2[2]]['TYPE'] = 4
			ant_arr1[e1[2]]['MIDX'] = e2[2]
			ant_arr2[e2[2]]['MIDX'] = e1[2]			
			
		elif opassed == True and icpassed == False:
			# mark as good orientation and bad insert size.
			ant_arr1[e1[2]]['TYPE'] = 5
			ant_arr2[e2[2]]['TYPE'] = 5
			ant_arr1[e1[2]]['MIDX'] = e2[2]
			ant_arr2[e2[2]]['MIDX'] = e1[2]			
			
		elif opassed == True and icpassed == True and ebpassed == False:
			# bad intra contig bound.
			ant_arr1[e1[2]]['TYPE'] = 6
			ant_arr2[e2[2]]['TYPE'] = 6
			ant_arr1[e1[2]]['MIDX'] = e2[2]
			ant_arr2[e2[2]]['MIDX'] = e1[2]		
		'''
		if opassed == False or icpassed == False or ebpassed == False:
			# write to ok.
			if oki >= ham_arr1_ok.size - 1:
				# increase by a million.
				ham_arr1_ok.resize((ham_arr1_ok.size + 1000000, 1))
				ham_arr2_ok.resize((ham_arr1_ok.size + 1000000, 1))
			
			# check if in this frame or last.
			if e1[0] == chunk_idx_1:
				ham_arr1_ok[oki]['QNAME'] = ham_arr1[e1[1]]['QNAME']
				ham_arr1_ok[oki]['RNAME'] = ham_arr1[e1[1]]['RNAME']
				ham_arr1_ok[oki]['FLAG'] = ham_arr1[e1[1]]['FLAG']
				ham_arr1_ok[oki]['POS'] = ham_arr1[e1[1]]['POS']
			else:
				ham_arr1_ok[oki]['QNAME'] = prev_ham_arr1[e1[1]]['QNAME']
				ham_arr1_ok[oki]['RNAME'] = prev_ham_arr1[e1[1]]['RNAME']
				ham_arr1_ok[oki]['FLAG'] = prev_ham_arr1[e1[1]]['FLAG']
				ham_arr1_ok[oki]['POS'] = prev_ham_arr1[e1[1]]['POS']				
			
			# check if in this frame or last.
			if e2[0] == chunk_idx_2:
				ham_arr2_ok[oki]['QNAME'] = ham_arr2[e2[1]]['QNAME']
				ham_arr2_ok[oki]['RNAME'] = ham_arr2[e2[1]]['RNAME']
				ham_arr2_ok[oki]['FLAG'] = ham_arr2[e2[1]]['FLAG']
				ham_arr2_ok[oki]['POS'] = ham_arr2[e2[1]]['POS']
			else:
				ham_arr2_ok[oki]['QNAME'] = prev_ham_arr2[e2[1]]['QNAME']
				ham_arr2_ok[oki]['RNAME'] = prev_ham_arr2[e2[1]]['RNAME']
				ham_arr2_ok[oki]['FLAG'] = prev_ham_arr2[e2[1]]['FLAG']
				ham_arr2_ok[oki]['POS'] = prev_ham_arr2[e2[1]]['POS']				
			oki += 1
		
		else:
			# its actually a good read!
			#ant_arr1[e1[2]]['TYPE'] = 1
			#ant_arr2[e2[2]]['TYPE'] = 1
			#ant_arr1[e1[2]]['MIDX'] = e2[2]
			#ant_arr2[e2[2]]['MIDX'] = e1[2]		
			
		#	print ham_arr1[e1[1]], ham_arr2[e2[1]]
			
			if peri == ham_arr1_perfect.size - 1:
				# increase by a million.
				ham_arr1_perfect.resize((ham_arr1_perfect.size + 1000000, 1))
				ham_arr2_perfect.resize((ham_arr1_perfect.size + 1000000, 1))
				print "Resized"
			
			if e1[0] == chunk_idx_1:
				ham_arr1_perfect[peri]['QNAME'] = ham_arr1[e1[1]]['QNAME']
				ham_arr1_perfect[peri]['RNAME'] = ham_arr1[e1[1]]['RNAME']
				ham_arr1_perfect[peri]['FLAG'] = ham_arr1[e1[1]]['FLAG']
				ham_arr1_perfect[peri]['POS'] = ham_arr1[e1[1]]['POS']
			else:
				ham_arr1_perfect[peri]['QNAME'] = prev_ham_arr1[e1[1]]['QNAME']
				ham_arr1_perfect[peri]['RNAME'] = prev_ham_arr1[e1[1]]['RNAME']
				ham_arr1_perfect[peri]['FLAG'] = prev_ham_arr1[e1[1]]['FLAG']
				ham_arr1_perfect[peri]['POS'] = prev_ham_arr1[e1[1]]['POS']
				
			if e2[0] == chunk_idx_2:
				ham_arr2_perfect[peri]['QNAME'] = ham_arr2[e2[1]]['QNAME']
				ham_arr2_perfect[peri]['RNAME'] = ham_arr2[e2[1]]['RNAME']
				ham_arr2_perfect[peri]['FLAG'] = ham_arr2[e2[1]]['FLAG']
				ham_arr2_perfect[peri]['POS'] = ham_arr2[e2[1]]['POS']
			else:

				ham_arr2_perfect[peri]['QNAME'] = prev_ham_arr2[e2[1]]['QNAME']
				ham_arr2_perfect[peri]['RNAME'] = prev_ham_arr2[e2[1]]['RNAME']
				ham_arr2_perfect[peri]['FLAG'] = prev_ham_arr2[e2[1]]['FLAG']
				ham_arr2_perfect[peri]['POS'] = prev_ham_arr2[e2[1]]['POS']

			peri += 1
			
		
	# 3: mult mapped	/	unmapped
	elif bsz1 > 1 and bsz2 == 0:
		for e in buf1:
			ant_arr1[e[2]]['TYPE'] = 2
		logging.debug(3)

	# 3: unmapped	/	mult mapped
	elif bsz1 == 0 and bsz2 > 1:
		for e in buf2:
			ant_arr2[e[2]]['TYPE'] = 2
		logging.debug(3)
		
	# 4: mult mapped	/	unique mapped
	elif bsz1 > 1 and bsz2 == 1:
		e2 = buf2.pop()
		for e1 in buf1:
			ant_arr1[e1[2]]['TYPE'] = 2
			ant_arr1[e1[2]]['MIDX'] = e2[2]
			
		ant_arr2[e2[2]]['TYPE'] = 1
		ant_arr2[e2[2]]['MIDX'] = e1[2]
		logging.debug(4)

	# 4: unique mapped	/	mult mapped
	elif bsz1 == 1 and bsz2 > 1:
		e1 = buf1.pop()
		for e2 in buf2:
			ant_arr2[e2[2]]['TYPE'] = 2
			ant_arr2[e2[2]]['MIDX'] = e1[2]
			
		ant_arr1[e1[2]]['TYPE'] = 1
		ant_arr1[e1[2]]['MIDX'] = e2[2]
		logging.debug(4)

	# 5: mult mapped	/	multi mapped
	elif bsz1 > 1 and bsz2 > 1:
		e2 = buf2.pop()
		for e1 in buf1:
			ant_arr1[e1[2]]['TYPE'] = 2
			ant_arr1[e1[2]]['MIDX'] = e2[2]
		buf2.add(e2)
		
		e1 = buf1.pop()
		for e2 in buf2:
			ant_arr2[e2[2]]['TYPE'] = 2
			ant_arr2[e2[2]]['MIDX'] = e1[2]

		logging.debug(5)
		
	else:
		print "fixme"
		print buf1
		print buf2
		sys.exit(1)
	
	
	# update the buffer.
	buf1 = set()
	buf2 = set()
	if choice == True:
		bufname = pname1
		buf1.add((chunk_idx_1, local_i, i))
	else:
		bufname = pname2
		buf2.add((chunk_idx_2, local_j, j))



# save annotation.
'''
logging.info("saving annotation.")
ant_obj = ham.HamAnnotate(ham_out_file, read=0)
ant_obj.save_annotation(ant_arr1, ant_arr2)
ant_obj.close()
'''
logging.info("saving annotation1.")
ant_obj = ham.HamFile(ham_out_file + "_1_ok.ham", read=0, cprs=True)
ant_obj.save_array(ham_arr1_ok)
ant_obj.close()

logging.info("saving annotation2.")
ant_obj2 = ham.HamFile(ham_out_file + "_1_perfect.ham", read=0, cprs=True)
ant_obj2.save_array(ham_arr1_perfect)
ant_obj2.close()

logging.info("saving annotation3.")
ant_obj3 = ham.HamFile(ham_out_file + "_2_ok.ham", read=0, cprs=True)
ant_obj3.save_array(ham_arr2_ok)
ant_obj3.close()

logging.info("saving annotation4.")
ant_obj4 = ham.HamFile(ham_out_file + "_2_perfect.ham", read=0, cprs=True)
ant_obj4.save_array(ham_arr2_perfect)
ant_obj4.close()

# Close pointer to ham object.
ham_obj1.close()
ham_obj2.close()

